import numpy as np
import matplotlib.pyplot as plt
import torch
import time


def main():
    # 示例数据集
    data = [
        [56769, '2015-04-17', '17:17:30', '6号线', '金桥路', '17:24:37', '6号线', '博兴路', 3.0],
        [56769, '2015-04-17', '19:59:57', '6号线', '博兴路', '20:04:07', '6号线', '金桥路', 3.0],
        # ...（其他数据行）
        [427031, '2015-04-17', '10:51:19', '6号线', '巨峰路', '11:03:06', '6号线', '金桥路', 2.7]
    ]

    def time_to_seconds(time_str):
        t = time.strptime(time_str, '%H:%M:%S')
        return t.tm_hour * 3600 + t.tm_min * 60 + t.tm_sec

    # 将时间转换为秒数
    for row in data:
        row[2] = time_to_seconds(row[2])
        row[5] = time_to_seconds(row[5])

    # 转换数据为矩阵
    data = np.array(data)[:, [0, 2, 5, 8]].astype(float)
    travel_time = data[:, 2] - data[:, 1]  # 计算出行时间

    # PSO算法参数
    num_particles = 150
    max_iter = 180
    dim = 3
    lb = -10 * np.ones(dim)
    ub = 10 * np.ones(dim)
    w_max = 0.9
    w_min = 0.4
    c1 = 2
    c2 = 2

    # 初始化粒子位置和速度
    pos = lb + (ub - lb) * np.random.rand(num_particles, dim)
    vel = np.zeros((num_particles, dim))

    # 初始化个体最优位置和全局最优位置
    pbest = pos.copy()
    pbest_val = np.zeros(num_particles)
    params = {'alpha': 1, 'beta': 1, 'gamma': 1, 'C_b': 1, 'H_min': 18000, 'H_max': 79200, 'C_k': 50}  # 示例参数

    for i in range(num_particles):
        pbest_val[i], Cx, Cy, Cv = objective_function(pbest[i], params, travel_time)
    gbest_idx = np.argmin(pbest_val)
    gbest = pbest[gbest_idx]
    gbest_val = pbest_val[gbest_idx]

    # 存储每次迭代的最优值
    best_values = np.zeros(max_iter)
    best_Cx = np.zeros(max_iter)
    best_Cy = np.zeros(max_iter)
    best_Cv = np.zeros(max_iter)

    # PSO主循环
    for iter in range(max_iter):
        w = w_max - (w_max - w_min) * (iter / max_iter)
        for i in range(num_particles):
            vel[i] = w * vel[i] + c1 * np.random.rand() * (pbest[i] - pos[i]) + c2 * np.random.rand() * (gbest - pos[i])
            pos[i] = pos[i] + vel[i] * 0.001
            pos[i] = np.maximum(np.minimum(pos[i], ub), lb)
            current_val, Cx, Cy, Cv = objective_function(pos[i], params, travel_time)
            if current_val < pbest_val[i]:
                pbest[i] = pos[i]
                pbest_val[i] = current_val
            if current_val < gbest_val:
                gbest = pos[i]
                gbest_val = current_val
                best_Cx[iter] = Cx
                best_Cy[iter] = Cy
                best_Cv[iter] = Cv
        best_values[iter] = gbest_val
        print(f'Iter {iter + 1}: Best Value = {gbest_val:.4f}, Cx = {best_Cx[iter]:.4f}, Cy = {best_Cy[iter]:.4f}, Cv = {best_Cv[iter]:.4f}')

    # 结果展示和图形化
    print('Optimization completed.')
    print(f'Best Value: {gbest_val}')
    print(f'Best Position: {gbest}')
    plt.figure()
    plt.subplot(4, 1, 1)
    plt.plot(range(1, max_iter + 1), best_values, '-o')
    plt.xlabel('Iteration')
    plt.ylabel('Best Value')
    plt.title('PSO Optimization Progress')
    plt.subplot(4, 1, 2)
    plt.plot(range(1, max_iter + 1), best_Cx, '-o')
    plt.xlabel('Iteration')
    plt.ylabel('C_x')
    plt.title('C_x Progress')
    plt.subplot(4, 1, 3)
    plt.plot(range(1, max_iter + 1), best_Cy, '-o')
    plt.xlabel('Iteration')
    plt.ylabel('C_y')
    plt.title('C_y Progress')
    plt.subplot(4, 1, 4)
    plt.plot(range(1, max_iter + 1), best_Cv, '-o')
    plt.xlabel('Iteration')
    plt.ylabel('C_v')
    plt.title('C_v Progress')
    plt.grid(True)
    plt.show()


# 目标函数定义
def objective_function(X, params, travel_time):
    alpha = params['alpha']
    beta = params['beta']
    gamma = params['gamma']
    C_b = params['C_b']
    H_min = params['H_min']
    H_max = params['H_max']
    C_k = params['C_k']

    # 从X向量中提取具体变量
    A_nk = X[:-2]
    T_Ai_k = X[-2]
    T_Li_k = X[-1]

    # 计算各部分成本
    C_d = abs(T_Ai_k - T_Li_k)
    Cx = C_b + C_d
    T_P_wait = abs(T_Ai_k - T_Li_k)
    Cy = C_b + T_P_wait
    Cv = T_P_wait

    # 添加约束条件
    penalty = 0

    # 约束1: 列车计划约束
    if T_Ai_k < H_min or T_Li_k > H_max:
        penalty += 1e6

    # 约束2: 列车容量限制约束
    if np.sum(A_nk) > C_k:
        penalty += 1e6

    # 约束3: 乘客服务约束
    if np.sum(A_nk) != 1:
        penalty += 1e6

    # 目标函数
    Z = alpha * np.sum(Cx * travel_time) + beta * np.sum(Cy * travel_time) + gamma * np.sum(Cv * travel_time) + penalty
    return Z, Cx, Cy, Cv


main()
